from twisted.internet import reactor
from twisted.names import server, dns, client
from twisted.python import log
from common_data import *
from datetime import datetime
from socket import gethostname, gethostbyname
import sys,os
import db_connection


class DNSRedirectServerFactory( server.DNSServerFactory ):

    current_config = None
    ip_address_list = None
    ip_address_to_rr_names = dict()
    ip_address_to_rr_details = dict()
    ip_address_to_environment_list = None
    
    config_last_modified = 0
    ip_address_list_last_modified = 0
    
    tcp_listener = None
    udp_listener = None
    
    local_hostname  = None
    
    def __init__(self, verbose = 0, standalone = False ):
        
        self.standalone = ( standalone and True ) or False
        
        # init db connection
        try:
            db_connection.connectToDb()
        except:
            log.err()
            
        # create default config if necessary    
        if not DNSServerConfig.select().count():
            DNSServerConfig()
        
        self.local_hostname = gethostname()
        
        # establish config loops  
        self.updateLatestConfigLoop()
        self.updateLatestIpsLoop()
        
        # establish initial resolvers
        if self.current_config[ 'resolvers' ]:
            resolver = client.Resolver( servers = self.current_config[ 'resolvers' ] )
        else:
            resolver = client.Resolver( resolv = '/etc/resolv.conf' )
    
        server.DNSServerFactory.__init__(self, verbose = verbose, clients = [ resolver ] )
        
        self.protocol = dns.DNSProtocol( self )
        self.datagram_protocol = dns.DNSDatagramProtocol( self )
        
        self.tcp_listener = self.assignTCPListener()
        self.udp_listener = self.assignUDPListener()
        
    def assignTCPListener(self):
        """ creates and returns tcp listener
        """
        if not self.current_config[ 'tcp_port' ]:
            raise Exception( 'cannot assign tcp listener: tcp_port not defined' )
        
        if not self.protocol:
            raise Exception( 'cannot assign tcp listener: protocol not defined' )
        
        if self.standalone:
            return reactor.listenTCP( self.current_config[ 'tcp_port' ], self )
        else:
            from twisted.application import internet
            return internet.TCPServer( self.current_config[ 'tcp_port' ], self )

    def assignUDPListener(self):
        """ creates and returns udp listener
        """
        if not self.current_config[ 'udp_port' ]:
            raise Exception( 'cannot assign udp listener: udp_port not defined' )
        
        if not self.datagram_protocol:
            raise Exception( 'cannot assign udp listener: datagram_protocol not defined' )
        
        if self.standalone:
            return reactor.listenUDP( self.current_config[ 'udp_port' ], self.datagram_protocol )
        else:
            from twisted.application import internet
            return internet.UDPServer( self.current_config[ 'udp_port' ], self.datagram_protocol )
            
    def updateLatestConfigLoop(self):
        """ check for newer config; if present loads the config and restarts services if necessary; then registers its next iteration
        """
        
        # check for newer config       
        try:
            newer_config = DNSServerConfig.select( DNSServerConfig.q.modified_time > self.config_last_modified ).count()
        except:
            log.err()
            sys.exit()
        
        # load newer config if present  
        if newer_config:
            log.msg( 'newer config found' )
            try:
                new_config = list( DNSServerConfig.select( orderBy=DNSServerConfig.q.modified_time ).reversed() )[0]
            except:
                log.err('cannot load server config')
                log.err()
            
            # determine if we are loading the initial config or if we are overwriting an existing config
            if not self.current_config:
                self.current_config = dict()
                self.current_config['resolvers'] = ( new_config.resolvers and list( ( resolver.hostname, resolver.port ) for resolver in new_config.resolvers ) ) or list()
                self.current_config['udp_port'] = new_config.udp_port
                self.current_config['tcp_port'] = new_config.tcp_port
                self.current_config['config_refresh_interval'] = new_config.config_refresh_interval
                self.current_config['active_list_refresh_interval'] = new_config.active_list_refresh_interval
                self.current_config['deny_inactive_ips'] = new_config.deny_inactive_ips
                self.current_config['auth_name'] = ( new_config.auth_name and new_config.auth_name ) or ""
                
                self.config_last_modified = datetime.now()
            else:
                if self.current_config[ 'udp_port' ] != new_config.udp_port:
                    self.current_config[ 'udp_port' ] = new_config.udp_port
                    
                    # restart listener on new udp port                    
                    if self.udp_listener:
                        if self.standalone:
                            self.udp_listener.stopListening().addCallback( 
                                lambda *x: setattr( self, 'udp_listener', self.assignUDPListener() ) )
                        else:
                            self.udp_listener.stopService().addCallback( 
                                lambda *x: setattr( self, 'udp_listener', self.assignUDPListener() ) ).addCallback( 
                                lambda *x: self.udp_listener.startService() )
                                            
                if self.current_config[ 'tcp_port' ] != new_config.tcp_port:
                    self.current_config[ 'tcp_port' ] = new_config.tcp_port
 
                    # restart listener on new tcp port
                    if self.tcp_listener:
                        if self.standalone:
                            self.tcp_listener.stopListening().addCallback(
                                lambda *x: setattr( self, 'tcp_listener', self.assignTCPListener() ) )
                        else:
                            self.tcp_listener.stopService().addCallback( 
                                lambda *x: setattr( self, 'tcp_listener', self.assignTCPListener() ) ).addCallback( 
                                lambda *x: self.tcp_listener.startService() )
                                                                        
                self.current_config['config_refresh_interval'] = new_config.config_refresh_interval
                self.current_config['active_list_refresh_interval'] = new_config.active_list_refresh_interval
                self.current_config['deny_inactive_ips'] = new_config.deny_inactive_ips
                
                if new_config.resolvers:
                    self.current_config['resolvers'] = list( ( resolver.hostname, resolver.port ) for resolver in new_config.resolvers )
                    resolver = client.Resolver( servers = self.current_config[ 'resolvers' ] )
                else:
                    self.current_config['resolvers'] = ()
                    resolver = client.Resolver( resolv = '/etc/resolv.conf' )
                
                self.resolver.resolvers = [ resolver ]
                
                self.current_config['auth_name'] = new_config.auth_name
                
                    
            self.config_last_modified = datetime.now()
            
        reactor.callLater( self.current_config[ 'config_refresh_interval' ], self.updateLatestConfigLoop )
            
    def updateLatestIpsLoop(self):    
        """  updates list of active ips; registers its next iteration 
        """

        # check for modified checkout, ip, or rr
        checkout_check = Checkout.select( Checkout.q.modified_time > self.ip_address_list_last_modified )
        ip_check = User.select( User.q.ip_address_modified_time > self.ip_address_list_last_modified )
        rr_check = ResourceRecord.select( ResourceRecord.q.modified_time > self.ip_address_list_last_modified )
        
        try:
            if checkout_check.count() or ip_check.count() or rr_check.count():
                
                # get ips
                ip_list = list( user.ip_address for user in list( User.select( User.q.ip_address != None ) ) )
                
                # get checkouts
                checkout_list = list( Checkout.select( orderBy = Checkout.q.id ) )
                
                self.ip_address_list = []
                self.ip_address_to_rr_names = {}
                self.ip_address_to_rr_details = {}
                
                # iterate through checkouts and match users to their ips;  generate active ip list
                for checkout in checkout_list:
                    rrs = checkout.environment.redirected_rrs
                    users = ( checkout.party.users + checkout.party.admins )
                    
                    for user in users:
                        if user.ip_address in ip_list:
                            ip = user.ip_address
                            
                            self.ip_address_list.append( ip )
                            
                            if ip in self.ip_address_to_rr_names:
                                self.ip_address_to_rr_names[ ip ].extend( list( rr.name for rr in rrs ) )
                                self.ip_address_to_rr_details[ ip ].update( dict( ( rr.name, rr ) for rr in rrs ) )   
                            else:
                                self.ip_address_to_rr_names[ ip ] = list( rr.name for rr in rrs )
                                self.ip_address_to_rr_details[ ip ] = dict( ( rr.name, rr ) for rr in rrs )
                                
                self.ip_address_list_last_modified = datetime.now()
        except:
            log.err('cannot retrieve ip address list count')
            log.err()
                
        reactor.callLater( self.current_config[ 'active_list_refresh_interval' ], self.updateLatestIpsLoop )
            
    def allowQuery(self, message, protocol, address):
        """ used to validate the ips of incoming query requests
        """
        
        # allow continue if were set to deny inactive clients
        if self.current_config['deny_inactive_ips']:
            ( address_ip, address_port ) = address
            if not address_ip in self.ip_address_list:
                log.msg( 'query denied: address_ip = %s' % address_ip )
                return False
        
        return server.DNSServerFactory.allowQuery( self, message, protocol, address)

    def handleQuery(self, message, protocol, address):
        """ parses incoming queries and redirects those that match any active checkout
        """
        
        query_name = str( message.queries[0].name )
        ( address_ip, address_port ) = address
        
        # redirect if query is from an active ip and is concerning an active checkout
        if address_ip in self.ip_address_list and query_name in self.ip_address_to_rr_names[ address_ip ]:
            log.msg( 'redirecting query: client = %s, query_name = %s' % ( address_ip, query_name ) )
            answer_name     = self.ip_address_to_rr_details[ address_ip ][ query_name ].name
            answer_ttl      = self.ip_address_to_rr_details[ address_ip ][ query_name ].ttl
            answer_address  = self.ip_address_to_rr_details[ address_ip ][ query_name ].value
            answer_type     = dns.REV_TYPES[ self.ip_address_to_rr_details[ address_ip ][ query_name ].type ]
            answer_klass    = dns.REV_CLASSES[ self.ip_address_to_rr_details[ address_ip ][ query_name ].klass ]
            auth_name       = self.current_config.get( 'auth_name', None ) or self.local_hostname
            
            answer = dns.RRHeader( name=answer_name, type=answer_type, cls=answer_klass, ttl=answer_ttl, payload=dns.Record_A( address=answer_address, ttl=answer_ttl ) )
            auth = dns.RRHeader( name=auth_name, type=dns.NS, ttl=answer_ttl, payload=dns.Record_NS( name=auth_name, ttl=answer_ttl ) )
            message.rCode = dns.OK
            message.answers.append( answer )
            message.auth = 1
            message.authority.append( auth )
            self.sendReply(protocol, message, address)
        else:
            return server.DNSServerFactory.handleQuery(self, message, protocol, address) 